#include"fast_lio/frontend.h"

namespace fast_lio{

void Frontend::addFrame(shared_ptr<Frame> p_frame){
    lock_guard<mutex> lck(frame_mutex_);
    frame_deq_.push_back(p_frame);
}

bool Frontend::frontendHandle(){
    {
        lock_guard<mutex> lck(frame_mutex_);
        if(frame_deq_.empty())
            return false;
        p_current_frame_ = frame_deq_.front();
        frame_deq_.pop_front();
        if(p_current_frame_->imu_deq_.empty())
            return false;
    }
    ROS_DEBUG_STREAM("frame id: "<<p_current_frame_->frame_id_
                            <<" imu size: "<<p_current_frame_->imu_deq_.size()
                            <<" points size: "<<p_current_frame_->pc_surf_->points.size());
    switch (status_)
    {
        case FrontendStatus::INITING:
            init();
            break;
        case FrontendStatus::TRACKING_GOOD:
            ROS_DEBUG("Tracking good");
        case FrontendStatus::TRACKING_BAD:
        {
            std::chrono::system_clock::time_point t1_, t2_;
            t1_ = std::chrono::system_clock::now();
            frameRegistration();
            t2_ = std::chrono::system_clock::now();
            ROS_DEBUG_STREAM("cost time   "<<std::chrono::duration_cast<std::chrono::microseconds>( t2_-t1_ ).count()*0.001<<"ms");
            break;
        }
        case FrontendStatus::LOST:
                ROS_DEBUG("track lost");
                break;
        default:
            break;
    }
    p_last_frame_ = p_current_frame_;
    return true;
}

void Frontend::init(){
    p_current_frame_->setKeyFrame();
    p_map_->insertKeyFrame(p_current_frame_);
    p_imu_handle_->IMUInit(p_current_frame_);

    Eigen::Matrix<double,State::DIM_OF_STATE_,State::DIM_OF_STATE_> cov;
    cov.setZero();
    V3D covPos = p_config_->init_pos_std_.array().square();
    V3D covVel = p_config_->init_vel_std_.array().square();
    V3D covQ = p_config_->init_q_std_.array().square();
    V3D covBa = p_config_->init_ba_std_.array().square();
    V3D covBg = p_config_->init_bg_std_.array().square();
    cov.block<3, 3>(State::IDX_pos_, State::IDX_pos_) =
        covPos.asDiagonal();  // pos
    cov.block<3, 3>(State::IDX_vel_, State::IDX_vel_) =
        covVel.asDiagonal();  // vel
    cov.block<3, 3>(State::IDX_q_, State::IDX_q_) =
        covQ.asDiagonal();  // att
    cov.block<3, 3>(State::IDX_bg_, State::IDX_bg_) =
        covBg.asDiagonal();  // bg
    cov.block<3, 3>(State::IDX_ba_, State::IDX_ba_) =
        covBa.asDiagonal();  // ba
    cov.block<3, 3>(State::IDX_gra_, State::IDX_gra_) =
        V3D(0.0001,0.0001,0.0001).asDiagonal();  // gravity
    // cov.block<2, 2>(State::IDX_gra_, State::IDX_gra_) =
    //     Eigen::Vector2d(0.001,0.001).asDiagonal();  // gravity
    p_current_frame_->p_state_->cov_P_ = cov;

    status_ = FrontendStatus::TRACKING_GOOD;
}

void Frontend::frameRegistration(){
    p_current_frame_->setState(p_last_frame_->getState());
    predictBatch();
    p_kdtree_local_.reset(new pcl::KdTreeFLANN<PointType>);
    p_kdtree_local_->setInputCloud(p_map_->getLocalMap());

    linear_state_ = p_current_frame_->getState();
    Eigen::Matrix<double,State::DIM_OF_STATE_, 1> dif_state_vec,update_state_vec;
    bool is_converge = false;
    MXD Hk_,Kk_,P_inv_,IKH_,Pk_;
    VXD residual_;
    for (int iterCount = 0; iterCount < p_config_->max_iteraters_; iterCount++)
    {
        laserCloudOri_->clear();
        coeffSel_->clear();
        calSurfResidual();

        const unsigned int DIM_OF_MEAS = laserCloudOri_->points.size();
        residual_.resize(DIM_OF_MEAS);
        Hk_.resize(DIM_OF_MEAS, State::DIM_OF_STATE_);
        Hk_.setZero();
        for (int i = 0; i < DIM_OF_MEAS; ++i) {
            V3D P2xyz(laserCloudOri_->points[i].x, laserCloudOri_->points[i].y,
                    laserCloudOri_->points[i].z);
            V3D coff_xyz(coeffSel_->points[i].x, coeffSel_->points[i].y,
                        coeffSel_->points[i].z);
            residual_(i) = coeffSel_->points[i].intensity;

            Hk_.block<1, 3>(i, State::IDX_pos_) = coff_xyz.transpose();
            SO3 rot_temp = linear_state_.q_.toRotationMatrix();
            Hk_.block<1, 3>(i, State::IDX_q_) =
            coff_xyz.transpose() *
            -Sophus::SO3d::hat(linear_state_.q_.toRotationMatrix() * P2xyz)
            *rightJacobiansInv(rot_temp.log());
        }
        Kk_.resize(State::DIM_OF_STATE_, DIM_OF_MEAS);
        P_inv_.resize(State::DIM_OF_STATE_, State::DIM_OF_STATE_);
        P_inv_ = p_current_frame_->getState().cov_P_.inverse();
        Kk_=(Hk_.transpose()*Hk_+ pow(p_config_->lidar_std_,2)*P_inv_).inverse()*Hk_.transpose();

        linear_state_.boxMinus(p_current_frame_->getState(),dif_state_vec);
        update_state_vec =  Kk_ * (-residual_ + Hk_ * dif_state_vec) - dif_state_vec;
        // update_state_vec =  -Kk_ * residual_ ;
        // cout<<"resdual norm "<<residual_.norm()<<endl;

        // Divergence determination
        for (int i = 0; i < update_state_vec.size(); i++) {
            if (isnan(update_state_vec[i])) {
                update_state_vec[i] = 0;
                ROS_DEBUG("update_state_vec has nan");
            }
        }

        // Update the state
        linear_state_.boxPlus(update_state_vec, linear_state_);
        if (update_state_vec.norm() <= 1e-2) {
            break;
        }
    }
    // Update only one time
    IKH_ = Eigen::Matrix<double, State::DIM_OF_STATE_,  State::DIM_OF_STATE_>::Identity() - Kk_ * Hk_;
    // Pk_ = IKH_ * linear_state_.cov_P_;
    Pk_ = IKH_ * linear_state_.cov_P_*IKH_.transpose();
    enforceSymmetry(Pk_);
    linear_state_.cov_P_ = Pk_;

    linear_state_.boxMinus(p_current_frame_->getState(),dif_state_vec);
    p_current_frame_->setState(linear_state_);

    if(judgeKeyframe()){
        p_current_frame_->setKeyFrame();
        p_map_->insertKeyFrame(p_current_frame_);
    }
}

void Frontend::predictBatch(){
    p_imu_handle_->predictBatch(p_current_frame_);
}

void Frontend::calSurfResidual(){
    // cout<<"surf "<<endl;
    Eigen::Matrix4d T1 = Eigen::Matrix4d::Identity();
    T1.block<3, 3>(0, 0) = linear_state_.q_.toRotationMatrix();
    T1.block<3, 1>(0, 3) = linear_state_.pos_;

    #pragma omp parallel for num_threads(p_config_->numberOfCores_)
    for(int pt_idx =0; pt_idx<p_current_frame_->pc_surf_->points.size();pt_idx++){
        PointType pt_tranform;
        PointType coeff;
        vector<int> point_search_Idx;
        vector<float> point_search_dis;
        
        tranformPoint(p_current_frame_->pc_surf_->points[pt_idx], pt_tranform, T1);
        // 从上一帧中寻找最近邻点
        p_kdtree_local_->nearestKSearch(pt_tranform, 5, point_search_Idx, point_search_dis);
        Eigen::Matrix<float, 5, 3> matA0;
        Eigen::Matrix<float, 5, 1> matB0;
        Eigen::Vector3f matX0;

        matA0.setZero();
        matB0.fill(-1);
        matX0.setZero();
        // 确保搜索点足够近
        if(point_search_dis[4] < max_dis_plane_for_match_)
        {
            for (int j = 0; j < 5; j++) {
                matA0(j, 0) = p_map_->getLocalMap()->points[point_search_Idx[j]].x;
                matA0(j, 1) = p_map_->getLocalMap()->points[point_search_Idx[j]].y;
                matA0(j, 2) = p_map_->getLocalMap()->points[point_search_Idx[j]].z;
            }

            // 假设平面方程为ax+by+cz+1=0，这里就是求方程的系数abc，d=1
            matX0 = matA0.colPivHouseholderQr().solve(matB0);

            // 平面方程的系数，也是法向量的分量
            float pa = matX0(0, 0);
            float pb = matX0(1, 0);
            float pc = matX0(2, 0);
            float pd = 1;

            // 单位法向量
            float ps = sqrt(pa * pa + pb * pb + pc * pc);
            pa /= ps; pb /= ps; pc /= ps; pd /= ps;

            // 检查平面是否合格，如果5个点中有点到平面的距离超过0.2m，那么认为这些点太分散了，不构成平面
            bool planeValid = true;
            for (int j = 0; j < 5; j++) {
                if (fabs(pa * p_map_->getLocalMap()->points[point_search_Idx[j]].x +
                            pb * p_map_->getLocalMap()->points[point_search_Idx[j]].y +
                            pc * p_map_->getLocalMap()->points[point_search_Idx[j]].z + pd) > 0.2) {
                    planeValid = false;
                    break;
                }
            }

            if(planeValid)
            {
                Eigen::Vector3d target_line_a( p_map_->getLocalMap()->points[point_search_Idx[0]].x,
                                                        p_map_->getLocalMap()->points[point_search_Idx[0]].y,
                                                        p_map_->getLocalMap()->points[point_search_Idx[0]].z);
                Eigen::Vector3d target_line_b( p_map_->getLocalMap()->points[point_search_Idx[2]].x,
                                                        p_map_->getLocalMap()->points[point_search_Idx[2]].y,
                                                        p_map_->getLocalMap()->points[point_search_Idx[2]].z);
                Eigen::Vector3d target_line_c( p_map_->getLocalMap()->points[point_search_Idx[4]].x,
                                                        p_map_->getLocalMap()->points[point_search_Idx[4]].y,
                                                        p_map_->getLocalMap()->points[point_search_Idx[4]].z);
                Eigen::Vector3d m_unit_vec_ab,m_unit_vec_ac,m_unit_vec_normal,vec_ad;
                m_unit_vec_ab = target_line_b - target_line_a;
                m_unit_vec_ab = m_unit_vec_ab / m_unit_vec_ab.norm();

                m_unit_vec_ac = target_line_c - target_line_a;
                m_unit_vec_ac = m_unit_vec_ac / m_unit_vec_ac.norm();

                m_unit_vec_normal = m_unit_vec_ab.cross(m_unit_vec_ac);
                m_unit_vec_normal.normalize();

                Eigen::Vector3d pt_transfromed(pt_tranform.x, pt_tranform.y, pt_tranform.z);
                vec_ad = pt_transfromed - target_line_a;

                coeff.x = m_unit_vec_normal[0];
                coeff.y = m_unit_vec_normal[1];
                coeff.z = m_unit_vec_normal[2];
                coeff.intensity = vec_ad.transpose()* m_unit_vec_normal;

                // 距离越大，s越小，是个距离惩罚因子（权重）
                V3D point_raw(p_current_frame_->pc_surf_->points[pt_idx].x,
                                                p_current_frame_->pc_surf_->points[pt_idx].y,
                                                p_current_frame_->pc_surf_->points[pt_idx].z);
                float s = 1 - 0.9 * fabs(coeff.intensity)/point_raw.norm();

                // 点到直线的垂线段单位向量
                coeff.x =  coeff.x;
                coeff.y =  coeff.y;
                coeff.z =  coeff.z;
                // 点到直线距离
                coeff.intensity =  coeff.intensity;

                // cout<<"coeff "<<coeff<<endl;
                // cout<<"point raw "<<p_current_frame_->pc_surface_sample_->points[pt_idx]<<endl;
                if (s > 0.9) {
                    lock_guard<mutex> lck(coeff_mutex_);
                    coeffSel_->push_back(coeff);
                    laserCloudOri_->push_back(p_current_frame_->pc_surf_->points[pt_idx]);
                }
            }
        }
    }
}

bool  Frontend::judgeKeyframe(){
    double theta = SO3(p_current_frame_->p_state_->q_ * state_last_kf_.q_.inverse()).log().norm();
    if(theta<p_config_->kf_angle_threshold_&&
    (p_current_frame_->p_state_->pos_ - state_last_kf_.pos_).norm()<p_config_->kf_trans_threshold_){
        return false;
    }else{
        state_last_kf_ = *p_current_frame_->p_state_;
        return true;
    }
}

}
